#include <U8g2lib.h>
#include <Wire.h>

/* PN532 I2C Configuration */
#include <PN532_I2C.h>
#include <PN532.h>
PN532_I2C pn532i2c(Wire);
PN532 nfc(pn532i2c);

/* OLED Software I2C Configuration */
U8G2_SSD1306_128X64_NONAME_1_SW_I2C u8g2(U8G2_R0, /*clock=*/7, /*data=*/6, /*reset=*/U8X8_PIN_NONE);

// 矩阵键盘配置
const byte ROWS = 4;
const byte COLS = 4;
char keys[ROWS][COLS] = {
  {'1','2','3','A'},
  {'4','5','6','B'},
  {'7','8','9','C'},
  {'*','0','#','D'}
};
byte rowPins[ROWS] = {2, 3, 4, 5};
byte colPins[COLS] = {8, 9, 10, 11};

// 系统状态枚举
enum SystemState {
  DEFAULT_STATE,
  ENTER_PASSWORD,
  SELECT_CARD_NFC,
  ENTER_AMOUNT,
  DISPLAY_BALANCE
};

// 卡片信息结构体
struct Card {
  char uid[9];  // 存储UID字符串
  int balance;  // 卡片余额
};

// 全局变量
const int MAX_CARDS = 3;       // 最大支持3张卡片
Card cards[MAX_CARDS];         // 卡片存储数组
int cardCount = 0;             // 当前已注册卡片数
SystemState currentState = DEFAULT_STATE;
char currentUID[9] = "";       // 当前操作的卡片UID
char passwordInput[5] = "";    // 密码输入缓冲区
char amountInput[6] = "";      // 金额输入缓冲区
uint8_t inputPos = 0;          // 输入位置指针
uint32_t displayTimeout = 0;   // 显示超时计时

void setup() {
  Serial.begin(115200);

  // 初始化键盘引脚
  for(byte c=0; c<COLS; c++) pinMode(colPins[c], INPUT_PULLUP);
  for(byte r=0; r<ROWS; r++) {
    pinMode(rowPins[r], OUTPUT);
    digitalWrite(rowPins[r], HIGH);
  }

  // OLED初始化
  u8g2.begin();
  u8g2.setFont(u8g2_font_6x10_tr);
  displayScreen("System Ready", "Swipe Card");

  // NFC初始化
  Wire.begin();
  nfc.begin();
  nfc.SAMConfig();
}

void loop() {
  checkNFC();
  checkKeyboard();
  manageDisplay();
}

void checkNFC() {
  static uint32_t lastCheck = 0;
  if(millis() - lastCheck < 500) return;

  uint8_t uid[7];
  uint8_t uidLength;
  
  if(nfc.readPassiveTargetID(PN532_MIFARE_ISO14443A, uid, &uidLength)) {
    char scannedUID[9];
    snprintf(scannedUID, sizeof(scannedUID), "%02X%02X%02X%02X", uid[0], uid[1], uid[2], uid[3]);

    // 状态处理
    switch(currentState) {
      case DEFAULT_STATE:
        handleCardPayment(scannedUID);
        break;
        
      case SELECT_CARD_NFC:
        handleRechargeSelection(scannedUID);
        break;
    }
  }
  lastCheck = millis();
}

void handleCardPayment(char* uid) {
  // 查找或创建卡片记录
  int index = findCard(uid);
  if(index == -1) {
    if(cardCount < MAX_CARDS) {
      index = registerNewCard(uid);
    } else {
      displayScreen("Card Full", "Max 3 Cards");
      delay(2000);
      resetToDefault();
      return;
    }
  }

  // 扣费逻辑
  if(cards[index].balance >= 10) {
    cards[index].balance -= 10;
  }
  
  // 显示结果
  char line2[20];
  snprintf(line2, sizeof(line2), "Balance: %d", cards[index].balance);
  displayScreen("Payment Success", line2);
  currentState = DISPLAY_BALANCE;
  displayTimeout = millis();
}

void handleRechargeSelection(char* uid) {
  int index = findCard(uid);
  if(index != -1) {
    strcpy(currentUID, uid);
    currentState = ENTER_AMOUNT;
    inputPos = 0;
    memset(amountInput, 0, sizeof(amountInput));
    displayScreen("Enter Amount", "Then press #");
  } else {
    displayScreen("Card Not Found", "Swipe Again");
    delay(2000);
    currentState = SELECT_CARD_NFC;
    displayScreen("Select Card", "Swipe Target");
  }
}

void checkKeyboard() {
  char key = scanKey();
  if(key == '\0') return;

  switch(currentState) {
    case DEFAULT_STATE:
      if(key == 'A') {
        currentState = ENTER_PASSWORD;
        inputPos = 0;
        memset(passwordInput, 0, sizeof(passwordInput));
        displayScreen("Admin Mode", "Enter Password:");
      }
      break;

    case ENTER_PASSWORD:
      if(key == '#') {
        if(strcmp(passwordInput, "1234") == 0) {
          currentState = SELECT_CARD_NFC;
          displayScreen("Select Card", "Swipe Target");
        } else {
          displayScreen("Wrong Password!", "Try Again");
          delay(2000);
          resetToDefault();
        }
      } else if(isdigit(key) && inputPos < 4) {
        passwordInput[inputPos++] = key;
        char placeholder[5] = "";
        for(int i=0; i<inputPos; i++) placeholder[i] = '*';
        displayScreen("Password:", placeholder);
      }
      break;

    case ENTER_AMOUNT:
      if(key == '#') {
        int amount = atoi(amountInput);
        int index = findCard(currentUID);
        if(index != -1) {
          cards[index].balance += amount;
          char line2[20];
          snprintf(line2, sizeof(line2), "New: %d", cards[index].balance);
          displayScreen("Recharge OK!", line2);
          delay(2000);
          resetToDefault();
        }
      } else if(isdigit(key) && inputPos < 5) {
        amountInput[inputPos++] = key;
        displayScreen("Amount:", amountInput);
      }
      break;
  }
}

// 辅助函数
int findCard(char* uid) {
  for(int i=0; i<cardCount; i++) {
    if(strcmp(cards[i].uid, uid) == 0) {
      return i;
    }
  }
  return -1;
}

int registerNewCard(char* uid) {
  if(cardCount >= MAX_CARDS) return -1;
  strcpy(cards[cardCount].uid, uid);
  cards[cardCount].balance = 500; // 初始余额
  return cardCount++;
}

void manageDisplay() {
  if(currentState == DISPLAY_BALANCE && millis() - displayTimeout > 2000) {
    resetToDefault();
  }
}

void resetToDefault() {
  currentState = DEFAULT_STATE;
  memset(currentUID, 0, sizeof(currentUID));
  displayScreen("System Ready", "Swipe Card");
}

char scanKey() {
  // 保持原有键盘扫描逻辑不变
  for(byte r=0; r<ROWS; r++) {
    digitalWrite(rowPins[r], LOW);
    delayMicroseconds(500);
    for(byte c=0; c<COLS; c++) {
      if(digitalRead(colPins[c]) == LOW) {
        digitalWrite(rowPins[r], HIGH);
        delay(50);
        return keys[r][c];
      }
    }
    digitalWrite(rowPins[r], HIGH);
  }
  return '\0';
}

void displayScreen(const char* line1, const char* line2) {
  u8g2.firstPage();
  do {
    u8g2.drawStr(0, 12, line1);
    u8g2.drawStr(0, 26, line2);
  } while(u8g2.nextPage());
}
